import logging
import datetime
import time
import json
import os
from flask_socketio import SocketIO
from flask import Flask, request
from nmap import PortScanner
import pandas as pd
import openpyxl
import re
import sys
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from project.colored_logging import configure_logging, console_table

app = Flask(__name__)

#----------------------------------------------------------------------------------
# 日志配置
#----------------------------------------------------------------------------------

# 定义日志文件路径和名称
current_date = datetime.datetime.now().strftime('%Y-%m-%d')
monthly_folder = datetime.datetime.now().strftime('%Y-%m')
GPDIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
log_folder_path = os.path.join(GPDIR, 'logs', monthly_folder)
if not os.path.exists(log_folder_path):
    os.makedirs(log_folder_path)

log_file_name = os.path.join(log_folder_path, f'scan_{current_date}.log')
file_handler = logging.FileHandler(log_file_name, mode='a', encoding='utf-8')

custom_keyword_colors = {
    'error': 'RED',
    'warning': 'YELLOW',
    'critical': 'MAGENTA',
    'successfully': 'GREEN'
}

logger = configure_logging(keywords=['error', 'warning', 'critical', 'successfully'],
                        keyword_colors=custom_keyword_colors,
                        handlers= [file_handler])


results = []

# 初始化SocketIO
socketio = SocketIO(app)

def remove_special_characters(text):
    text = re.sub(r'[\x00-\x1F\x7F]', '', text)
    text = re.sub(r'\s', '', text)
    return text

def is_valid_ip(ip):
    parts = ip.split('.')
    if len(parts)!= 4:
        return False
    for part in parts:
        if not part.isdigit():
            return False
        num = int(part)
        if num < 0 or num > 255:
            return False
    return True

#----------------------------------------------------------------------------------
# 扫描主程序（扫描程序；数据排版；Excel数据写入）
#----------------------------------------------------------------------------------

@app.route('/scan', methods=['POST'])
def scan_flask():
    start_time = time.time()
    nm = PortScanner()
    results.clear()
    data = request.get_json()
    sum = 0
    for item in data:
        sum += 1
    data_received = "The data was received successfully."
    logging.info(data_received)
    socketio.emit('log_message', data_received)
    process_start = "The scanning process is starting."
    logging.info(f"The task begins. There are [{sum}] tasks in total..")
    logging.info(process_start)
    socketio.emit('log_message', process_start)

    # 开始扫描
    for index, item in enumerate(data, start=1):
        ip = item.get('ip')
        port = item.get('port')
        begin_message = f"[{index}/{sum}] Begin scanning【{ip}:{port}】"
        logging.info(begin_message)
        socketio.emit('log_message', begin_message)
        if ip is None or not is_valid_ip(ip):
            formatted_result = f"Host: {ip}, Port: {port}, State: Error, Service: Error"
            results.append(formatted_result)
            error_message = f"Error at index {index}: Invalid ip {ip}. Skipping this item."
            logging.error(error_message)
            socketio.emit('log_message', error_message)
            continue

        if port is None or port.lower() == 'none':
            # 当端口为None时进行全端口扫描，并记录相关信息
            warning_message = f"Port is None at index {index} for {ip}: Performing full port scan."
            logging.warning(warning_message)
            socketio.emit('log_message', warning_message)
            try:
                full_port_result = nm.scan(hosts=ip, arguments='-sS', timeout=1000)
                if full_port_result and 'scan' in full_port_result and full_port_result['scan']:
                    host = list(full_port_result['scan'].keys())[0]
                    for port_info_key, port_info in full_port_result['scan'][host]['tcp'].items():
                        formatted_result = f"Host: {host}, Port: {port_info_key}, State: {port_info['state']}, Service: {port_info['name']}"
                        results.append(formatted_result)
                else:
                    formatted_result = f"Host: {ip}, Port: {port}, State: Dead, Service: null"
                    warning_message = f"The scan result for the target {ip} is empty."
                    logging.warning(warning_message)
                    socketio.emit('log_message', warning_message)
                    results.append(formatted_result)
            except Exception as e:
                detailed_error_message = f"Error during full port scan of {ip}: {type(e).__name__}: {str(e)}"
                logging.error(detailed_error_message)
                socketio.emit('log_message', detailed_error_message)
        else:
            # 当端口不为None时进行单个端口扫描，并记录相关信息
            try:
                if not port.isdigit():
                    formatted_result = f"Host: {ip}, Port: {port}, State: Error, Service: Error"
                    results.append(formatted_result)
                    error_message = f"{ip}Port is not an integer.Skipping this item."
                    logging.error(error_message)
                    socketio.emit('log_message', error_message)
                    continue
                port_num = int(port)
                if port_num < 0 or port_num >= 65536:
                    formatted_result = f"Host: {ip}, Port: {port}, State: Error, Service: Error"
                    results.append(formatted_result)
                    error_message = f"{ip} Port is out of range (0-65535).Skipping this item."
                    logging.error(error_message)
                    socketio.emit('log_message', error_message)
                    continue
                single_port_result = nm.scan(hosts=ip, arguments=f'-sS -p {port}', timeout=1000)
                if 'scan' in single_port_result and single_port_result['scan']:
                    host = list(single_port_result['scan'].keys())[0]
                    port_info = single_port_result['scan'][host]['tcp'].get(int(port))
                    if port_info:
                        formatted_result = f"Host: {host}, Port: {port}, State: {port_info['state']}, Service: {port_info['name']}"
                        results.append(formatted_result)
                else:
                    second_scan_result = nm.scan(hosts=ip, arguments=f'-sT -Pn -p {port}', timeout=1000)
                    if 'scan' in second_scan_result and second_scan_result['scan']:
                        host = list(second_scan_result['scan'].keys())[0]
                        port_info = second_scan_result['scan'][host]['tcp'].get(int(port))
                        if port_info:
                            formatted_result = f"Host: {host}, Port: {port}, State: {port_info['state']}, Service: {port_info['name']}"
                            results.append(formatted_result)
                        else:
                            host = ip
                            warning_message = f"The scan result for the target {host} is empty after second scan."
                            logging.warning(warning_message)
                            socketio.emit('log_message', warning_message)
                            formatted_result = f"Host: {ip}, Port: {port}, State: closed, Service: {port_info['name']}"
                            results.append(formatted_result)
            except Exception as e:
                detailed_error_message = f"Error during scanning of {ip}:{port}: {type(e).__name__}: {str(e)}"
                logging.error(detailed_error_message)
                socketio.emit('log_message', detailed_error_message)
        scan_success_message = f"Scanning of {ip}:{port} completed successfully."
        logging.info(scan_success_message)
        socketio.emit('log_message', scan_success_message)
    end_time = time.time()
    scan_time_message = f"Scan took {end_time - start_time:.2f} seconds."

    # 对results进行规范化排版和排序
    formatted_results = []
    for result in results:
        parts = result.split(', ')
        host = parts[0].split(': ')[1]
        port = parts[1].split(': ')[1]
        state = parts[2].split(': ')[1]
        service = parts[3].split(': ')[1]

        # 设置列宽，同样假设列宽，可按需调整
        host_fmt = '{:<18}'.format(host)
        port_fmt = '{:<15}'.format(port)
        state_fmt = '{:<15}'.format(state)
        service_fmt = '{:<15}'.format(service)

        formatted_results.append((host_fmt, port_fmt, state_fmt, service_fmt))
    
    sorted_results = sorted(formatted_results, key=lambda x: (x[0], str(x[1]), x[2], x[3]))
    df = pd.DataFrame(sorted_results, columns=['{:<18}'.format('Host'), '{:<15}'.format('Port'), '{:<15}'.format('State'), '{:<15}'.format('Service')])
    df = df.reset_index(drop=True)

    success_message = (f"Scanning process completed successfully.\n{df.to_string(index=True)}")
    # logging.info(success_message)
    # socketio.emit('log_message', success_message)
    console_table(results,"Scanning process completed successfully.")
    logging.info(scan_time_message)
    socketio.emit('log_message', scan_time_message)

    # 将结果覆盖写入到Excel文件中
    data_folder_path = os.path.join(GPDIR, 'data')
    if not os.path.exists(data_folder_path):
        os.makedirs(data_folder_path)
    save_file_name = os.path.join(data_folder_path, 'nmap_results.xlsx')
    logging.info(f"Saving the scanning results separately.")
    socketio.emit('log_message', f"Saving the scanning results separately.")
    for _ in range(3):
        try:
            if os.path.exists(save_file_name):
                workbook = openpyxl.load_workbook(save_file_name)
                sheet = workbook.active
                # 清空现有内容进行覆盖写入
                sheet.delete_rows(2, sheet.max_row)
                for host, port, state, service in sorted_results:
                    sheet.append(list(map(remove_special_characters, [host, port, state, service])) )
                workbook.save(save_file_name)
                logging.info(f"Data overwritten to existing Excel file.")
                socketio.emit('log_message', f"Data overwritten to existing Excel file.")
                break
            else:
                workbook = openpyxl.Workbook()
                sheet = workbook.active
                sheet.append(['Host', 'Port', 'State', 'Service'])
                for host, port, state, service in sorted_results:
                    sheet.append(list(map(remove_special_characters, [host, port, state, service])) )
                workbook.save(save_file_name)
                logging.info(f"New Excel file created and data written.")
                socketio.emit('log_message', f"New Excel file created and data written.")
                break
        except PermissionError as e:
            logging.error("Error: Excel file occupied. Wait for 10s.")
            socketio.emit('log_message', "Error: Excel file occupied. Wait for 10s.")
            time.sleep(10)

    scan_finsh = f"The scanning program ended."
    logging.info(scan_finsh)
    socketio.emit('log_message', scan_finsh)

    return json.dumps({"results": results}), 200


@socketio.on('connect')
def handle_connect():
    logging.info('Client connected.')


@socketio.on('disconnect')
def handle_disconnect():
    logging.info(f'Client disconnected.')

if __name__ == '__main__':
    try:
        socketio.run(app, host='0.0.0.0', port=5001)
    finally:
        # 在程序结束时将日志缓冲区中的内容写入文件
        logging.shutdown()
